#  find the Barker code pair with the highest peak level and the smallest maximum absolute value of the sidelobe level.
import csv
import datetime
import math
import random
import barkerpair
from typing import List
import matplotlib.pyplot as plt


def fitness(barker_code_pair: List[List[int]]):
    """
    Evaluate fitness of a Barker code pair. Fitness is the peak level and the maximum absolute value of the
    sidelobe level.
    """
    N = len(
        barker_code_pair[0])  # suppose each individual code in pair has same length
    autocorr = [0]*N
    for shift in range(N):
        for i in range(N-shift):
            autocorr[shift] += barker_code_pair[0][i] * \
                barker_code_pair[1][i+shift]
        autocorr[shift] = abs(autocorr[shift])
    # sidelobel level ascending order, peak level descending order.
    return [max(autocorr[1:]), -autocorr[0]]


def generate_population(population_size: int, code_length: int):
    """
    Generate a population of Barker code pairs of length `code_length`.
    """
    population = []
    for _ in range(population_size):
        x = [random.choice([1, -1]) for _ in range(code_length)]
        y = [random.choice([1, -1]) for _ in range(code_length)]
        population.append([x, y])
    return population


def crossover1(parent1: List[List[int]], parent2: List[List[int]]):
    """
    Generate a new Barker code pair by crossover of two parent codes.
    """
    x = []
    y = []
    for i in range(len(parent1[0])):
        if random.random() < 0.5:
            x.append(parent1[0][i])
        else:
            x.append(parent2[0][i])

        if random.random() < 0.5:
            y.append(parent1[1][i])
        else:
            y.append(parent2[1][i])
    return [x, y]

def crossover2(parent1: List[List[int]], parent2: List[List[int]]):
    """
    Generate a new Barker code pair by two-point crossover of two parent codes.
    """
    x = []
    y = []
    length = len(parent1[0])
    point1 = random.randint(0, length - 1)
    point2 = random.randint(0, length - 1)
    if point1 > point2:
        point1, point2 = point2, point1
    for i in range(length):
        if i < point1 or i > point2:
            x.append(parent1[0][i])
            y.append(parent1[1][i])
        else:
            x.append(parent2[0][i])
            y.append(parent2[1][i])
    return [x, y]

def crossover(parent1: List[List[int]], parent2: List[List[int]]):
    if random.random() < 0.5:
        return crossover1(parent1,parent2)
    else:
        return crossover2(parent1,parent2)

# def mutate(code: List[List[int]], mutation_rate: float):
#     """
#     Mutate a Barker code pair by flipping a random bit with probability `mutation_rate`.
#     """
#     for i in range(len(code[0])):
#         if random.random() < mutation_rate:
#             code[0][i] = -code[0][i]
#         if random.random() < mutation_rate:
#             code[1][i] = -code[1][i]
#     return code

def mutate(code: List[List[int]], mutation_rate: float, generation: int, generations: int):
    """
    Mutate a Barker code pair by flipping a random bit with probability `mutation_rate`.
    The mutation rate is varied sinusoidally over time.
    """
    sine_mutation_rate = (math.sin(generation) + 1) / 2
    for i in range(len(code[0])):
        if random.random() < sine_mutation_rate * mutation_rate:
            code[0][i] = -code[0][i]
        if random.random() < sine_mutation_rate * mutation_rate:
            code[1][i] = -code[1][i]
    return code

def tournament_selection(population: List[List[List[int]]], tournament_size: int):
    """
    Select an individual for reproduction using tournament selection.
    """
    tournament = random.sample(population, tournament_size)
    winner = max(tournament, key=lambda x: fitness(x))
    return winner

def roulette_wheel_selection(population: List[List[List[int]]]):
    """
    Select an individual for reproduction using roulette wheel selection.
    """
    fitnesses = [fitness(x) for x in population]
    total_fitness = sum(fitnesses)
    probabilities = [f/total_fitness for f in fitnesses]
    cumulative_probabilities = [sum(probabilities[:i+1]) for i in range(len(probabilities))]
    pointer = random.random()
    for i in range(len(cumulative_probabilities)):
        if pointer <= cumulative_probabilities[i]:
            return population[i]
        

def genetic_algorithm(population_size: int, code_length: int, mutation_rate: float, generations: int):
    """
    Use genetic algorithm to find Barker code pairs of length `code_length`.
    """
    population = generate_population(population_size, code_length)
    best_fitnesses = []
    tournament_size = 3
    for generation in range(generations):
        population = sorted(population, key=lambda x: fitness(x))
        fittest = population[0]
        best_fitness = fitness(fittest)
        best_fitness[1] = -best_fitness[1]
        best_fitnesses.append(best_fitness)
        print("Generation", generation+1, "Fitness:", best_fitness)

        next_generation = population[0:30]
        while len(next_generation) < population_size:
            parent1 = random.choice(population[:population_size//2])
            parent2 = random.choice(population[:population_size//2])
            child = crossover(parent1, parent2)
            # child = mutate(child, mutation_rate)
            child = mutate(child, mutation_rate, generation, generations)
            next_generation.append(child)
        population = next_generation
    plt.plot(range(1, generations+1), [bf[0]
             for bf in best_fitnesses], label="sidelobe level")
    plt.plot(range(1, generations+1), [bf[1]
             for bf in best_fitnesses], label="Peak level")
    plt.xlabel("Generation")
    plt.ylabel("Fitness Function")
    plt.legend()
    plt.savefig(f"fitness_plot_{code_length}.png")
    plt.show()
    return population


if __name__ == '__main__':
    population_size = 1000
    code_length = 21
    mutation_rate = 0.02
    generations = 100000
    best_code = genetic_algorithm(
        population_size, code_length, mutation_rate, generations)
    now = datetime.datetime.now()
    time_stamp = now.strftime("%Y%m%d%H%M%S")
    filename = f'barkerpair-{code_length}-{time_stamp}.csv'
    with open(filename, 'w') as f:
        writer = csv.writer(f)
        writer.writerow(["length", "seqpair", "autocorrelation",
                             "main peak/sidelobe", "main peak/length"])
        for barker_code_pair in best_code:
            autoCorr = barkerpair.get_autocorrelation(barker_code_pair)
            if barkerpair.is_barker_pair(autoCorr):
                writer.writerow([code_length, barker_code_pair, autoCorr, abs(autoCorr[0]), math.floor(abs(autoCorr[0]/code_length)*100)/100])
